rsmp4decrypt 0.1.0

//! Rust bindings and a native CLI for Bento4 `mp4decrypt`.
//!
//! The API accepts the same core inputs as the Bento4 tool:
//! decryption keys, an optional fragments info file, and optional progress reporting.

mod error;

pub use error::Error;

use core::ffi::{c_char, c_int, c_uchar, c_uint, c_void};
use std::{ffi::CString, path::Path, ptr, sync::Mutex};

pub const BENTO4_TAG: &str = "v1.6.0-641";

const KEY_KIND_TRACK_ID: c_uint = 0;
const KEY_KIND_KID: c_uint = 1;

type ProgressCallback =
    Option<unsafe extern "C" fn(step: c_uint, total: c_uint, user_data: *mut c_void)>;

// Bento4 decryption is not reliable when multiple contexts run concurrently in one process.
static AP4_LOCK: Mutex<()> = Mutex::new(());

#[repr(C)]
struct FfiKeyEntry {
    kind: c_uint,
    track_id: c_uint,
    kid: [u8; 16],
    key: [u8; 16],
}

unsafe extern "C" {
    fn rsmp4decrypt_context_new(keys: *const FfiKeyEntry, key_count: c_uint) -> *mut c_void;
    fn rsmp4decrypt_context_free(ctx: *mut c_void);
    fn rsmp4decrypt_free(ptr: *mut c_uchar);
    fn rsmp4decrypt_decrypt_file(
        ctx: *mut c_void,
        input_path: *const c_char,
        output_path: *const c_char,
        fragments_info_path: *const c_char,
        progress: ProgressCallback,
        user_data: *mut c_void,
    ) -> c_int;
    fn rsmp4decrypt_decrypt_memory(
        ctx: *mut c_void,
        input_data: *const c_uchar,
        input_size: c_uint,
        fragments_info_data: *const c_uchar,
        fragments_info_size: c_uint,
        progress: ProgressCallback,
        user_data: *mut c_void,
        output_data: *mut *mut c_uchar,
        output_size: *mut c_uint,
    ) -> c_int;
}

/// Identifies a decryption key either by track ID or by KID.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub enum KeyId {
    TrackId(u32),
    Kid([u8; 16]),
}

/// A single Bento4 decryption key entry.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct DecryptionKey {
    id: KeyId,
    key: [u8; 16],
}

impl DecryptionKey {
    /// Parses the same key syntax accepted by the CLI: `<id>:<key>`.
    pub fn from_spec(spec: &str) -> Result<Self, Error> {
        let (id_text, key_text) = spec.split_once(':').ok_or_else(|| Error::InvalidKeySpec {
            input: spec.to_owned(),
            message: "expected <id>:<key>".to_owned(),
        })?;

        let key = parse_hex_16(key_text)?;
        let id = if id_text.len() == 32 {
            KeyId::Kid(parse_hex_16(id_text)?)
        } else {
            let track_id = id_text.parse::<u32>().map_err(|_| Error::InvalidKeySpec {
                input: spec.to_owned(),
                message: "track ids must be unsigned decimal integers".to_owned(),
            })?;
            KeyId::TrackId(track_id)
        };

        Ok(Self { id, key })
    }

    /// Creates a key entry addressed by a 128-bit KID.
    pub fn kid(kid: &str, key: &str) -> Result<Self, Error> {
        Ok(Self {
            id: KeyId::Kid(parse_hex_16(kid)?),
            key: parse_hex_16(key)?,
        })
    }

    /// Creates a key entry addressed by a track ID.
    pub fn track(track_id: u32, key: &str) -> Result<Self, Error> {
        Ok(Self {
            id: KeyId::TrackId(track_id),
            key: parse_hex_16(key)?,
        })
    }

    /// Returns the identifier used to select this key during decryption.
    pub fn id(&self) -> KeyId {
        self.id
    }

    /// Returns the raw 16-byte content key.
    pub fn key_bytes(&self) -> [u8; 16] {
        self.key
    }

    fn to_ffi(self) -> FfiKeyEntry {
        match self.id {
            KeyId::TrackId(track_id) => FfiKeyEntry {
                kind: KEY_KIND_TRACK_ID,
                track_id,
                kid: [0; 16],
                key: self.key,
            },
            KeyId::Kid(kid) => FfiKeyEntry {
                kind: KEY_KIND_KID,
                track_id: 0,
                kid,
                key: self.key,
            },
        }
    }
}

/// Builder for [`Mp4Decryptor`].
#[derive(Debug, Default)]
pub struct Mp4DecryptorBuilder {
    keys: Vec<DecryptionKey>,
}

impl Mp4DecryptorBuilder {
    /// Adds a pre-parsed decryption key.
    pub fn key(mut self, key: DecryptionKey) -> Self {
        self.keys.push(key);
        self
    }

    /// Adds a key entry using the CLI `<id>:<key>` format.
    pub fn key_spec(self, spec: &str) -> Result<Self, Error> {
        Ok(self.key(DecryptionKey::from_spec(spec)?))
    }

    /// Adds a key addressed by a 128-bit KID.
    pub fn kid_key(self, kid: &str, key: &str) -> Result<Self, Error> {
        Ok(self.key(DecryptionKey::kid(kid, key)?))
    }

    /// Adds a key addressed by a track ID.
    pub fn track_key(self, track_id: u32, key: &str) -> Result<Self, Error> {
        Ok(self.key(DecryptionKey::track(track_id, key)?))
    }

    /// Builds a decryptor instance.
    pub fn build(self) -> Result<Mp4Decryptor, Error> {
        if self.keys.is_empty() {
            return Err(Error::NoKeys);
        }

        let ffi_keys = self
            .keys
            .into_iter()
            .map(DecryptionKey::to_ffi)
            .collect::<Vec<_>>();

        let ptr = unsafe { rsmp4decrypt_context_new(ffi_keys.as_ptr(), ffi_keys.len() as c_uint) };
        if ptr.is_null() {
            return Err(Error::ContextCreationFailed);
        }

        Ok(Mp4Decryptor {
            ptr,
            _keys: ffi_keys,
        })
    }
}

/// A reusable Bento4 decryptor context.
pub struct Mp4Decryptor {
    ptr: *mut c_void,
    _keys: Vec<FfiKeyEntry>,
}

unsafe impl Send for Mp4Decryptor {}
unsafe impl Sync for Mp4Decryptor {}

impl Mp4Decryptor {
    /// Creates a builder for configuring decryption keys.
    pub fn builder() -> Mp4DecryptorBuilder {
        Mp4DecryptorBuilder::default()
    }

    /// Decrypts MP4 bytes that are already in memory.
    ///
    /// When decrypting a fragmented segment, pass the matching init segment as `fragments_info`.
    pub fn decrypt<I, F>(&self, input_data: I, fragments_info: Option<F>) -> Result<Vec<u8>, Error>
    where
        I: AsRef<[u8]>,
        F: AsRef<[u8]>,
    {
        self.decrypt_impl(input_data, fragments_info, None, ptr::null_mut())
    }

    /// Decrypts MP4 bytes in memory and reports progress through a callback.
    pub fn decrypt_with_progress<I, F, P>(
        &self,
        input_data: I,
        fragments_info: Option<F>,
        progress: &mut P,
    ) -> Result<Vec<u8>, Error>
    where
        I: AsRef<[u8]>,
        F: AsRef<[u8]>,
        P: FnMut(u32, u32),
    {
        with_progress_callback(progress, |callback, user_data| {
            self.decrypt_impl(input_data, fragments_info, callback, user_data)
        })
    }

    /// Decrypts an MP4 file on disk.
    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>,
    {
        self.decrypt_file_impl(
            input_path,
            output_path,
            fragments_info_path,
            None,
            ptr::null_mut(),
        )
    }

    /// Decrypts an MP4 file on disk and reports progress through a callback.
    pub fn decrypt_file_with_progress<I, O, F, P>(
        &self,
        input_path: I,
        output_path: O,
        fragments_info_path: Option<F>,
        progress: &mut P,
    ) -> Result<(), Error>
    where
        I: AsRef<Path>,
        O: AsRef<Path>,
        F: AsRef<Path>,
        P: FnMut(u32, u32),
    {
        with_progress_callback(progress, |callback, user_data| {
            self.decrypt_file_impl(
                input_path,
                output_path,
                fragments_info_path,
                callback,
                user_data,
            )
        })
    }

    fn decrypt_impl<I, F>(
        &self,
        input_data: I,
        fragments_info: Option<F>,
        progress: ProgressCallback,
        user_data: *mut c_void,
    ) -> Result<Vec<u8>, Error>
    where
        I: AsRef<[u8]>,
        F: AsRef<[u8]>,
    {
        let input_data = input_data.as_ref();
        let input_size =
            u32::try_from(input_data.len()).map_err(|_| Error::DataTooLarge(u32::MAX))?;

        let (fragments_ptr, fragments_size) = match fragments_info.as_ref() {
            Some(data) => {
                let data = data.as_ref();
                let size = u32::try_from(data.len()).map_err(|_| Error::DataTooLarge(u32::MAX))?;
                (data.as_ptr(), size)
            }
            None => (ptr::null(), 0),
        };

        let mut output_ptr = ptr::null_mut();
        let mut output_size = 0;

        let result = {
            let _guard = AP4_LOCK.lock().expect("poisoned Bento4 lock");
            unsafe {
                rsmp4decrypt_decrypt_memory(
                    self.ptr,
                    input_data.as_ptr(),
                    input_size,
                    fragments_ptr,
                    fragments_size,
                    progress,
                    user_data,
                    &mut output_ptr,
                    &mut output_size,
                )
            }
        };

        if result != 0 {
            return Err(Error::DecryptionFailed(result));
        }

        let bytes = unsafe {
            if output_ptr.is_null() {
                Vec::new()
            } else {
                let slice = std::slice::from_raw_parts(output_ptr, output_size as usize);
                let output = slice.to_vec();
                rsmp4decrypt_free(output_ptr);
                output
            }
        };

        Ok(bytes)
    }

    fn decrypt_file_impl<I, O, F>(
        &self,
        input_path: I,
        output_path: O,
        fragments_info_path: Option<F>,
        progress: ProgressCallback,
        user_data: *mut c_void,
    ) -> Result<(), Error>
    where
        I: AsRef<Path>,
        O: AsRef<Path>,
        F: AsRef<Path>,
    {
        let input_path = path_to_cstring(input_path.as_ref())?;
        let output_path = path_to_cstring(output_path.as_ref())?;
        let fragments_info_path = match fragments_info_path {
            Some(path) => Some(path_to_cstring(path.as_ref())?),
            None => None,
        };

        let result = {
            let _guard = AP4_LOCK.lock().expect("poisoned Bento4 lock");
            unsafe {
                rsmp4decrypt_decrypt_file(
                    self.ptr,
                    input_path.as_ptr(),
                    output_path.as_ptr(),
                    fragments_info_path
                        .as_ref()
                        .map_or(ptr::null(), |path| path.as_ptr()),
                    progress,
                    user_data,
                )
            }
        };

        if result == 0 {
            Ok(())
        } else {
            Err(Error::DecryptionFailed(result))
        }
    }
}

impl Drop for Mp4Decryptor {
    fn drop(&mut self) {
        let _guard = AP4_LOCK.lock().expect("poisoned Bento4 lock");
        unsafe { rsmp4decrypt_context_free(self.ptr) }
    }
}

fn parse_hex_16(input: &str) -> Result<[u8; 16], Error> {
    let bytes = hex::decode(input)?;
    if bytes.len() != 16 {
        return Err(Error::InvalidHex {
            input: input.to_owned(),
            message: format!("expected 16 bytes, got {}", bytes.len()),
        });
    }

    let mut output = [0; 16];
    output.copy_from_slice(&bytes);
    Ok(output)
}

fn path_to_cstring(path: &Path) -> Result<CString, Error> {
    CString::new(path.to_string_lossy().as_bytes()).map_err(Error::from)
}

struct ProgressThunk<P> {
    callback: *mut P,
    invoke: unsafe fn(*mut P, u32, u32),
}

unsafe extern "C" fn progress_trampoline<P>(step: c_uint, total: c_uint, user_data: *mut c_void)
where
    P: FnMut(u32, u32),
{
    let thunk = unsafe { &mut *(user_data.cast::<ProgressThunk<P>>()) };
    unsafe { (thunk.invoke)(thunk.callback, step, total) };
}

unsafe fn invoke_progress<P>(callback: *mut P, step: u32, total: u32)
where
    P: FnMut(u32, u32),
{
    let callback = unsafe { &mut *callback };
    callback(step, total);
}

fn with_progress_callback<P, R>(
    progress: &mut P,
    run: impl FnOnce(ProgressCallback, *mut c_void) -> R,
) -> R
where
    P: FnMut(u32, u32),
{
    let mut thunk = ProgressThunk {
        callback: progress as *mut P,
        invoke: invoke_progress::<P>,
    };

    run(
        Some(progress_trampoline::<P>),
        (&mut thunk as *mut ProgressThunk<P>).cast::<c_void>(),
    )
}