hydra

shim.rs (15712B)

---

 //! Safe Rust wrapper over the Objective-C tap shim (`tap_shim.m`).
 //!
 //! [`MonitorRoute`] owns a live monitor: a process tap feeding a private aggregate
 //! device whose IOProc copies the tapped audio to a hardware output. Dropping it tears
 //! the CoreAudio objects down in the correct order.
 
 use std::ffi::{c_void, CString};
 use std::os::raw::c_char;
 use std::path::PathBuf;
 use std::ptr;
 use std::sync::atomic::{AtomicBool, AtomicI32, AtomicU32, AtomicU64, Ordering};
 use std::sync::Arc;
 
 use anyhow::{bail, Result};
 
 use super::process::process_object_for_pid;
 use super::AudioObjectID;
 
 /// Shared realtime parameters. Layout MUST match `HydraParams` in `tap_shim.m` exactly —
 /// C `_Atomic int`/`unsigned int`/`unsigned long long` map to `AtomicI32`/`AtomicU32`/
 /// `AtomicU64` (same size/align), and `float*` to `*mut f32`. A runtime size assertion in
 /// `MonitorRoute::start_combined` guards against drift.
 #[repr(C)]
 struct HydraParams {
     gain: f32,
     muted: i32,
     peak: [f32; 8],
     running: i32,
     callbacks: u64,
     // recording ring (SPSC)
     rec_on: AtomicI32,
     rec_buf: *mut f32,
     rec_cap: u32,
     rec_channels: u32,
     rec_write: AtomicU64,
     rec_read: AtomicU64,
     rec_overruns: AtomicU64,
     // tap format, published by the IOProc / start
     fmt_channels: AtomicU32,
     fmt_sample_rate: AtomicU32,
 }
 
 /// Opaque-to-Rust handle the shim fills in. Layout must match `HydraRoute`.
 #[repr(C)]
 struct HydraRoute {
     tap: AudioObjectID,
     aggregate: AudioObjectID,
     ioproc: *mut c_void,
     params: *mut HydraParams,
 }
 
 extern "C" {
     fn hydra_monitor_start(
         proc_objs: *const AudioObjectID,
         n_procs: i32,
         output_uid: *const c_char,
         params: *mut HydraParams,
         out_route: *mut HydraRoute,
     ) -> i32;
     fn hydra_monitor_stop(route: *mut HydraRoute);
     /// `sizeof(HydraParams)` as the C compiler sees it — for the layout assertion.
     fn hydra_params_size() -> usize;
     /// Route a hardware INPUT device (e.g. the MacBook mic) to an output. No tap/TCC.
     fn hydra_input_start(
         input_uid: *const c_char,
         output_uid: *const c_char,
         params: *mut HydraParams,
         out_route: *mut HydraRoute,
     ) -> i32;
 }
 
 static NEXT_ID: AtomicU64 = AtomicU64::new(1);
 
 /// A running monitor route. Tearing down on `Drop` is what keeps coreaudiod from wedging.
 pub struct MonitorRoute {
     id: String,
     /// The PIDs whose audio this route taps and mixes together (the "combine" feature).
     pids: Vec<i32>,
     raw: HydraRoute,
     /// Heap-stable params the C IOProc reads; we keep it alive for the route's lifetime.
     params: Box<HydraParams>,
     /// Active recording, if any. Always stopped (thread joined) before `params` drops.
     recorder: Option<Recorder>,
 }
 
 /// A live recording: owns the ring storage the IOProc writes into and the drain thread that
 /// flushes it to a WAV file. The IOProc holds a raw pointer into `ring`/`params`; this struct
 /// guarantees the drain thread is joined and `rec_on` cleared before either is freed.
 struct Recorder {
     path: PathBuf,
     stop: Arc<AtomicBool>,
     handle: Option<std::thread::JoinHandle<std::io::Result<u64>>>,
     /// Kept alive so the raw `rec_buf` pointer the IOProc holds stays valid.
     _ring: Box<[f32]>,
 }
 
 /// Wrap a raw params pointer so it can cross into the drain thread. Safe because the params
 /// box outlives the thread (joined in `Recorder::stop`) and we only touch atomics + the ring.
 struct ParamsPtr(*const HydraParams);
 unsafe impl Send for ParamsPtr {}
 
 // The CoreAudio handles are just IDs; the only pointer is into our own boxed params,
 // which outlives any access. Safe to move between threads (the daemon holds it behind a Mutex).
 unsafe impl Send for MonitorRoute {}
 
 /// Allocate a fresh, heap-stable `HydraParams` for a new route. Asserts the Rust/C layouts
 /// match (the struct is shared with the realtime IOProc; a mismatch would corrupt audio-thread
 /// memory) before any route is created.
 fn new_params(gain: f32) -> Box<HydraParams> {
     assert_eq!(
         std::mem::size_of::<HydraParams>(),
         unsafe { hydra_params_size() },
         "HydraParams layout mismatch between Rust and tap_shim.m"
     );
     Box::new(HydraParams {
         gain,
         muted: 0,
         peak: [0.0; 8],
         running: 0,
         callbacks: 0,
         rec_on: AtomicI32::new(0),
         rec_buf: ptr::null_mut(),
         rec_cap: 0,
         rec_channels: 0,
         rec_write: AtomicU64::new(0),
         rec_read: AtomicU64::new(0),
         rec_overruns: AtomicU64::new(0),
         fmt_channels: AtomicU32::new(0),
         fmt_sample_rate: AtomicU32::new(0),
     })
 }
 
 impl MonitorRoute {
     /// Tap a single `pid` and monitor it to `output_uid` (or the default output if `None`).
     pub fn start(pid: i32, output_uid: Option<&str>, gain: f32) -> Result<Self> {
         Self::start_combined(&[pid], output_uid, gain)
     }
 
     /// Tap several PIDs at once, mixing their audio into one route. The CATapDescription is
     /// built over every process's tap object (`initStereoMixdownOfProcesses:`), so the
     /// aggregate's input is the sum of all of them.
     pub fn start_combined(pids: &[i32], output_uid: Option<&str>, gain: f32) -> Result<Self> {
         if pids.is_empty() {
             bail!("a route needs at least one source");
         }
         let proc_objs: Vec<AudioObjectID> =
             pids.iter().map(|&p| process_object_for_pid(p)).collect::<Result<_>>()?;
 
         let mut params = new_params(gain);
         let mut raw = HydraRoute { tap: 0, aggregate: 0, ioproc: ptr::null_mut(), params: ptr::null_mut() };
 
         let c_uid = output_uid.map(CString::new).transpose()?;
         let uid_ptr = c_uid.as_ref().map_or(ptr::null(), |c| c.as_ptr());
 
         let st = unsafe {
             hydra_monitor_start(
                 proc_objs.as_ptr(),
                 proc_objs.len() as i32,
                 uid_ptr,
                 params.as_mut() as *mut _,
                 &mut raw,
             )
         };
         if st != 0 {
             bail!("starting monitor for pids {pids:?} failed: {}", os_status(st));
         }
 
         Ok(Self {
             id: format!("r{}", NEXT_ID.fetch_add(1, Ordering::Relaxed)),
             pids: pids.to_vec(),
             raw,
             params,
             recorder: None,
         })
     }
 
     /// Route a hardware input device (by UID — e.g. the MacBook mic) to `output_uid` (or the
     /// default output). No process tap, no TCC consent: a plain input→output aggregate.
     pub fn start_input(input_uid: &str, output_uid: Option<&str>, gain: f32) -> Result<Self> {
         let mut params = new_params(gain);
         let mut raw = HydraRoute { tap: 0, aggregate: 0, ioproc: ptr::null_mut(), params: ptr::null_mut() };
 
         let c_in = CString::new(input_uid)?;
         let c_out = output_uid.map(CString::new).transpose()?;
         let out_ptr = c_out.as_ref().map_or(ptr::null(), |c| c.as_ptr());
 
         let st = unsafe { hydra_input_start(c_in.as_ptr(), out_ptr, params.as_mut() as *mut _, &mut raw) };
         if st != 0 {
             bail!("starting input route for {input_uid} failed: {}", os_status(st));
         }
 
         Ok(Self {
             id: format!("r{}", NEXT_ID.fetch_add(1, Ordering::Relaxed)),
             pids: Vec::new(), // hardware input has no PID
             raw,
             params,
             recorder: None,
         })
     }
 
     pub fn id(&self) -> &str {
         &self.id
     }
 
     pub fn pids(&self) -> &[i32] {
         &self.pids
     }
 
     pub fn set_gain(&mut self, gain: f32) {
         unsafe { ptr::write_volatile(&mut self.params.gain, gain) };
     }
 
     pub fn set_muted(&mut self, muted: bool) {
         unsafe { ptr::write_volatile(&mut self.params.muted, muted as i32) };
     }
 
     pub fn gain(&self) -> f32 {
         unsafe { ptr::read_volatile(&self.params.gain) }
     }
 
     pub fn muted(&self) -> bool {
         unsafe { ptr::read_volatile(&self.params.muted) != 0 }
     }
 
     /// How many times the IOProc has fired (0 ⇒ the aggregate is never being clocked).
     pub fn callbacks(&self) -> u64 {
         unsafe { ptr::read_volatile(&self.params.callbacks) }
     }
 
     /// Most recent peak written by the IOProc (0.0..~1.0). Stored in `peak[0]`.
     pub fn peak(&self) -> f32 {
         unsafe { ptr::read_volatile(&self.params.peak[0]) }
     }
 
     pub fn is_recording(&self) -> bool {
         self.recorder.is_some()
     }
 
     /// The file currently being recorded to, if any.
     pub fn recording_path(&self) -> Option<&std::path::Path> {
         self.recorder.as_ref().map(|r| r.path.as_path())
     }
 
     /// Begin recording the tapped audio to a WAV file at `path`. The IOProc fills an SPSC
     /// ring; a drain thread writes the file. Errors if already recording or if the tap
     /// format isn't known yet (no callback has fired — start playback first).
     pub fn start_recording(&mut self, path: PathBuf) -> Result<()> {
         if self.recorder.is_some() {
             bail!("already recording");
         }
         let channels = self.params.fmt_channels.load(Ordering::Acquire);
         let sample_rate = self.params.fmt_sample_rate.load(Ordering::Acquire);
         if channels == 0 || sample_rate == 0 {
             bail!("tap format not ready yet — start audio in the app first, then record");
         }
 
         // ~4 seconds of ring headroom; the drain thread empties it continuously.
         let cap = (sample_rate as usize * channels as usize * 4).next_power_of_two();
         let mut ring = vec![0.0f32; cap].into_boxed_slice();
 
         // Point the IOProc at the ring and arm it. Order matters: set buf/cap/channels and
         // reset indices BEFORE flipping rec_on, so the audio thread never sees a half-armed ring.
         self.params.rec_buf = ring.as_mut_ptr();
         self.params.rec_cap = cap as u32;
         self.params.rec_channels = channels;
         self.params.rec_read.store(0, Ordering::Relaxed);
         self.params.rec_write.store(0, Ordering::Relaxed);
         self.params.rec_overruns.store(0, Ordering::Relaxed);
         self.params.rec_on.store(1, Ordering::Release);
 
         let stop = Arc::new(AtomicBool::new(false));
         let stop_t = Arc::clone(&stop);
         let pp = ParamsPtr(self.params.as_ref() as *const HydraParams);
         let out_path = path.clone();
         let handle = std::thread::spawn(move || drain_to_wav(pp, stop_t, &out_path, channels, sample_rate));
 
         self.recorder = Some(Recorder { path, stop, handle: Some(handle), _ring: ring });
         Ok(())
     }
 
     /// Stop recording: disarm the IOProc, join the drain thread, finalize the WAV. Returns
     /// the number of sample-frames written.
     pub fn stop_recording(&mut self) -> Result<u64> {
         let Some(mut rec) = self.recorder.take() else {
             bail!("not recording");
         };
         // Disarm the audio thread first so it stops touching the ring, then drain+join.
         self.params.rec_on.store(0, Ordering::Release);
         rec.stop.store(true, Ordering::Release);
         let frames = match rec.handle.take().map(|h| h.join()) {
             Some(Ok(Ok(floats))) => floats / self.params.rec_channels.max(1) as u64,
             Some(Ok(Err(e))) => bail!("recording write error: {e}"),
             Some(Err(_)) => bail!("recording thread panicked"),
             None => 0,
         };
         // ring (rec._ring) frees here, after the thread is joined and the IOProc disarmed.
         Ok(frames)
     }
 }
 
 impl Drop for MonitorRoute {
     fn drop(&mut self) {
         // Stop recording (joins the drain thread, disarms the IOProc) BEFORE tearing down
         // CoreAudio and freeing params — otherwise the audio thread could touch freed memory.
         if self.recorder.is_some() {
             let _ = self.stop_recording();
         }
         unsafe { hydra_monitor_stop(&mut self.raw) };
         // `params` Box is freed here, after the IOProc is guaranteed gone.
     }
 }
 
 /// Drain the SPSC ring to a 32-bit-float WAV until `stop` is set and the ring is empty.
 /// Runs on its own thread; the only shared state is atomics + the ring buffer.
 fn drain_to_wav(
     pp: ParamsPtr,
     stop: Arc<AtomicBool>,
     path: &std::path::Path,
     channels: u32,
     sample_rate: u32,
 ) -> std::io::Result<u64> {
     use std::io::Write;
     let params = unsafe { &*pp.0 };
     let cap = params.rec_cap as usize;
 
     let file = std::fs::File::create(path)?;
     let mut w = std::io::BufWriter::new(file);
     write_wav_header(&mut w, channels, sample_rate)?;
 
     let mut total_floats: u64 = 0;
     let mut scratch: Vec<f32> = Vec::with_capacity(cap);
     loop {
         let write = params.rec_write.load(Ordering::Acquire);
         let read = params.rec_read.load(Ordering::Relaxed);
         let avail = (write - read) as usize;
         if avail > 0 {
             scratch.clear();
             for i in 0..avail {
                 scratch.push(params.rec_buf_get((read as usize + i) % cap));
             }
             let mut bytes = Vec::with_capacity(avail * 4);
             for &s in &scratch {
                 bytes.extend_from_slice(&s.to_le_bytes());
             }
             w.write_all(&bytes)?;
             params.rec_read.store(write, Ordering::Release);
             total_floats += avail as u64;
         } else if stop.load(Ordering::Acquire) {
             break;
         } else {
             std::thread::sleep(std::time::Duration::from_millis(20));
         }
     }
     w.flush()?;
     finalize_wav(w.into_inner().map_err(|e| e.into_error())?, total_floats)?;
     Ok(total_floats)
 }
 
 impl HydraParams {
     /// Read one float from the recording ring (drain thread only).
     fn rec_buf_get(&self, idx: usize) -> f32 {
         unsafe { *self.rec_buf.add(idx) }
     }
 }
 
 /// Write a canonical 44-byte WAV header for 32-bit IEEE float PCM, with placeholder sizes
 /// (patched by [`finalize_wav`] once the total is known).
 fn write_wav_header<W: std::io::Write>(w: &mut W, channels: u32, sample_rate: u32) -> std::io::Result<()> {
     let ch = channels as u16;
     let byte_rate = sample_rate * channels * 4;
     let block_align = (channels * 4) as u16;
     w.write_all(b"RIFF")?;
     w.write_all(&0u32.to_le_bytes())?; // RIFF size — patched later
     w.write_all(b"WAVE")?;
     w.write_all(b"fmt ")?;
     w.write_all(&16u32.to_le_bytes())?; // fmt chunk size
     w.write_all(&3u16.to_le_bytes())?; // format 3 = IEEE float
     w.write_all(&ch.to_le_bytes())?;
     w.write_all(&sample_rate.to_le_bytes())?;
     w.write_all(&byte_rate.to_le_bytes())?;
     w.write_all(&block_align.to_le_bytes())?;
     w.write_all(&32u16.to_le_bytes())?; // bits per sample
     w.write_all(b"data")?;
     w.write_all(&0u32.to_le_bytes())?; // data size — patched later
     Ok(())
 }
 
 /// Patch the RIFF + data sizes in a finished WAV now that the float count is known.
 fn finalize_wav(mut file: std::fs::File, total_floats: u64) -> std::io::Result<()> {
     use std::io::{Seek, SeekFrom, Write};
     let data_bytes = (total_floats * 4) as u32;
     file.seek(SeekFrom::Start(4))?;
     file.write_all(&(36 + data_bytes).to_le_bytes())?; // RIFF size = 36 + data
     file.seek(SeekFrom::Start(40))?;
     file.write_all(&data_bytes.to_le_bytes())?; // data chunk size
     file.flush()
 }
 
 /// Render an OSStatus as a four-char-code if printable, else a number — CoreAudio errors
 /// are usually FourCCs (e.g. `!pri` = no permission, `!dev` = bad device).
 fn os_status(st: i32) -> String {
     let b = (st as u32).to_be_bytes();
     if b.iter().all(|&c| c.is_ascii_graphic() || c == b' ') {
         format!("'{}{}{}{}' ({st})", b[0] as char, b[1] as char, b[2] as char, b[3] as char)
     } else {
         st.to_string()
     }
 }