Skip to main content

pcmflux/
ogg_sink.rs

1/*
2 * This Source Code Form is subject to the terms of the Mozilla Public
3 * License, v. 2.0. If a copy of the MPL was not distributed with this
4 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
5 */
6
7//! Ogg Opus fan-out on a Unix socket: every consumer that connects receives the stream's
8//! `OpusHead` and `OpusTags` pages, then one page per encoded packet, so the socket reads
9//! as a standard Ogg Opus stream (`ffmpeg -i unix:<path>`, the pixelflux recorder).
10//!
11//! The tap never perturbs the capture: the capture thread only hands each page to a
12//! bounded per-consumer channel drained by that consumer's writer thread, and a consumer
13//! that falls behind is dropped rather than waited on.
14
15use std::fs;
16use std::io::{ErrorKind, Write};
17use std::os::unix::net::UnixListener;
18use std::sync::atomic::{AtomicBool, Ordering};
19use std::sync::mpsc::{SyncSender, TrySendError, sync_channel};
20use std::sync::{Arc, Mutex};
21use std::thread;
22use std::time::Duration;
23
24const WRITE_TIMEOUT: Duration = Duration::from_millis(100);
25const ACCEPT_POLL_INTERVAL: Duration = Duration::from_millis(50);
26/// Per-consumer backlog: two seconds of 20 ms packets before a consumer is dropped.
27const CLIENT_QUEUE_CAP: usize = 100;
28/// The Ogg page checksum's polynomial, applied to the whole page with the field zeroed.
29const CRC_POLY: u32 = 0x04c1_1db7;
30
31/// What `OpusHead` tells a decoder; `mapping` is `(streams, coupled, table)` for channel
32/// mapping family 1, absent for family 0 (mono and stereo).
33pub struct OpusHead {
34    pub channels: u8,
35    pub pre_skip: u16,
36    pub input_sample_rate: u32,
37    pub mapping: Option<(u8, u8, Vec<u8>)>,
38}
39
40struct Client {
41    tx: SyncSender<Arc<Vec<u8>>>,
42    stop: Arc<AtomicBool>,
43}
44
45/// One Ogg Opus stream served to every consumer of a Unix socket.
46pub struct OggSink {
47    path: String,
48    clients: Arc<Mutex<Vec<Client>>>,
49    shutdown: Arc<AtomicBool>,
50    serial: u32,
51    seq: u32,
52}
53
54fn crc32(data: &[u8]) -> u32 {
55    let mut crc = 0u32;
56    for &b in data {
57        crc ^= (b as u32) << 24;
58        for _ in 0..8 {
59            crc = if crc & 0x8000_0000 != 0 {
60                (crc << 1) ^ CRC_POLY
61            } else {
62                crc << 1
63            };
64        }
65    }
66    crc
67}
68
69/// One page carrying one packet: the lacing table splits it into 255-byte segments, a
70/// final short segment (empty when the length is a multiple of 255) closing it.
71fn page(serial: u32, seq: u32, granule: u64, header_type: u8, packet: &[u8]) -> Vec<u8> {
72    let mut lacing: Vec<u8> = vec![255; packet.len() / 255];
73    lacing.push((packet.len() % 255) as u8);
74    let mut out = Vec::with_capacity(27 + lacing.len() + packet.len());
75    out.extend_from_slice(b"OggS");
76    out.push(0);
77    out.push(header_type);
78    out.extend_from_slice(&granule.to_le_bytes());
79    out.extend_from_slice(&serial.to_le_bytes());
80    out.extend_from_slice(&seq.to_le_bytes());
81    out.extend_from_slice(&[0u8; 4]);
82    out.push(lacing.len() as u8);
83    out.extend_from_slice(&lacing);
84    out.extend_from_slice(packet);
85    let crc = crc32(&out).to_le_bytes();
86    out[22..26].copy_from_slice(&crc);
87    out
88}
89
90/// The identification and comment headers as the stream's first two pages.
91fn header_pages(serial: u32, head: &OpusHead) -> Vec<u8> {
92    let mut id = b"OpusHead".to_vec();
93    id.push(1);
94    id.push(head.channels);
95    id.extend_from_slice(&head.pre_skip.to_le_bytes());
96    id.extend_from_slice(&head.input_sample_rate.to_le_bytes());
97    id.extend_from_slice(&0i16.to_le_bytes());
98    match &head.mapping {
99        None => id.push(0),
100        Some((streams, coupled, table)) => {
101            id.push(1);
102            id.push(*streams);
103            id.push(*coupled);
104            id.extend_from_slice(table);
105        }
106    }
107    let mut tags = b"OpusTags".to_vec();
108    tags.extend_from_slice(&(b"pcmflux".len() as u32).to_le_bytes());
109    tags.extend_from_slice(b"pcmflux");
110    tags.extend_from_slice(&0u32.to_le_bytes());
111    let mut out = page(serial, 0, 0, 0x02, &id);
112    out.extend_from_slice(&page(serial, 1, 0, 0, &tags));
113    out
114}
115
116impl OggSink {
117    /// Bind the socket at `path`, or say why it cannot be bound. A file there that another
118    /// account owns is left alone and refused: in a shared directory it is that account's
119    /// listener, and a reader of this stream would reach it.
120    pub fn bind(path: &str, head: &OpusHead) -> Result<Self, String> {
121        use std::os::unix::fs::MetadataExt;
122        let stale = match fs::symlink_metadata(path) {
123            Ok(meta) if meta.uid() != unsafe { libc::geteuid() } => {
124                Err("another account owns it".to_string())
125            }
126            Ok(_) => fs::remove_file(path).map_err(|e| e.to_string()),
127            Err(_) => Ok(()),
128        };
129        let listener = stale
130            .and_then(|_| {
131                UnixListener::bind(path)
132                    .and_then(|l| l.set_nonblocking(true).map(|_| l))
133                    .map_err(|e| e.to_string())
134            })
135            .map_err(|e| format!("output_socket {path} cannot be bound: {e}"))?;
136        let serial = std::process::id()
137            ^ (std::time::SystemTime::now()
138                .duration_since(std::time::UNIX_EPOCH)
139                .map(|d| d.subsec_nanos())
140                .unwrap_or(0));
141        let headers = Arc::new(header_pages(serial, head));
142        let clients: Arc<Mutex<Vec<Client>>> = Arc::new(Mutex::new(Vec::new()));
143        let shutdown = Arc::new(AtomicBool::new(false));
144        let (clients_acc, shutdown_acc, path_log) =
145            (clients.clone(), shutdown.clone(), path.to_string());
146        thread::spawn(move || {
147            while !shutdown_acc.load(Ordering::Relaxed) {
148                match listener.accept() {
149                    Ok((stream, _)) => {
150                        let _ = stream.set_write_timeout(Some(WRITE_TIMEOUT));
151                        let (tx, rx) = sync_channel::<Arc<Vec<u8>>>(CLIENT_QUEUE_CAP);
152                        let stop = Arc::new(AtomicBool::new(false));
153                        let stop_writer = stop.clone();
154                        let _ = tx.try_send(headers.clone());
155                        thread::spawn(move || {
156                            let mut stream = stream;
157                            for bytes in rx.iter() {
158                                if stop_writer.load(Ordering::Relaxed)
159                                    || stream.write_all(&bytes).is_err()
160                                {
161                                    break;
162                                }
163                            }
164                        });
165                        clients_acc.lock().unwrap().push(Client { tx, stop });
166                        eprintln!("[pcmflux] ogg sink consumer connected on {path_log}");
167                    }
168                    Err(e) if e.kind() == ErrorKind::WouldBlock => {
169                        thread::sleep(ACCEPT_POLL_INTERVAL)
170                    }
171                    Err(e) => {
172                        eprintln!("[pcmflux] ogg sink accept error: {e}");
173                        thread::sleep(Duration::from_millis(500));
174                    }
175                }
176            }
177        });
178        Ok(Self {
179            path: path.to_string(),
180            clients,
181            shutdown,
182            serial,
183            seq: 2,
184        })
185    }
186
187    /// Send one packet as a page; `granule` is the 48 kHz sample count at its end.
188    pub fn write_packet(&mut self, packet: &[u8], granule: u64) {
189        let bytes = Arc::new(page(self.serial, self.seq, granule, 0, packet));
190        self.seq = self.seq.wrapping_add(1);
191        let mut clients = self.clients.lock().unwrap();
192        clients.retain(|c| match c.tx.try_send(bytes.clone()) {
193            Ok(()) => true,
194            Err(TrySendError::Full(_)) | Err(TrySendError::Disconnected(_)) => {
195                c.stop.store(true, Ordering::Relaxed);
196                eprintln!("[pcmflux] ogg sink dropping a consumer that stopped reading");
197                false
198            }
199        });
200    }
201}
202
203impl Drop for OggSink {
204    fn drop(&mut self) {
205        self.shutdown.store(true, Ordering::Relaxed);
206        for c in self.clients.lock().unwrap().drain(..) {
207            c.stop.store(true, Ordering::Relaxed);
208        }
209        let _ = fs::remove_file(&self.path);
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216    use std::io::Read;
217    use std::os::unix::net::UnixStream;
218
219    /// A leftover this account owns is removed and the socket bound, and the sink removes
220    /// its socket when it is dropped.
221    #[test]
222    fn a_stale_socket_of_this_session_is_replaced() {
223        let path = format!("/tmp/pcmflux-ogg-stale-{}.sock", std::process::id());
224        fs::write(&path, b"stale").unwrap();
225        let head = OpusHead {
226            channels: 2,
227            pre_skip: 312,
228            input_sample_rate: 48000,
229            mapping: None,
230        };
231        let sink = OggSink::bind(&path, &head);
232        assert!(
233            sink.is_ok(),
234            "an own leftover is removed and the socket bound"
235        );
236        drop(sink);
237        assert!(
238            fs::symlink_metadata(&path).is_err(),
239            "the socket is removed with the sink"
240        );
241    }
242
243    /// A path that cannot be bound is an error naming it, not a sink that serves nobody.
244    #[test]
245    fn an_unbindable_path_is_an_error() {
246        let path = format!("/tmp/pcmflux-ogg-missing-{}/rec.sock", std::process::id());
247        let head = OpusHead {
248            channels: 2,
249            pre_skip: 312,
250            input_sample_rate: 48000,
251            mapping: None,
252        };
253        let err = OggSink::bind(&path, &head)
254            .err()
255            .expect("a missing directory is refused");
256        assert!(err.contains(&path), "{err}");
257    }
258
259    /// A consumer that connects reads the two header pages and then every packet page, each
260    /// with the granule it was given. A packet that fills its last segment exactly is closed by
261    /// a zero lacing value, and the checksum is Ogg's own, whose check value over `123456789` is
262    /// 0x89a1897f: recomputing a page's checksum with this module's function alone would pass
263    /// for any CRC-32, while libogg rejects every page of another.
264    #[test]
265    fn stream_pages_are_well_formed() {
266        assert_eq!(crc32(b"123456789"), 0x89a1_897f);
267        let path = format!("/tmp/pcmflux-ogg-test-{}.sock", std::process::id());
268        let head = OpusHead {
269            channels: 2,
270            pre_skip: 312,
271            input_sample_rate: 48000,
272            mapping: None,
273        };
274        let mut sink = OggSink::bind(&path, &head).expect("bind");
275        let mut consumer = UnixStream::connect(&path).expect("connect");
276        consumer
277            .set_read_timeout(Some(Duration::from_secs(2)))
278            .unwrap();
279        thread::sleep(Duration::from_millis(150));
280        sink.write_packet(&[0xfc, 1, 2, 3], 960);
281        sink.write_packet(&vec![0xfc; 255], 1920);
282        thread::sleep(Duration::from_millis(150));
283        let mut buf = vec![0u8; 4096];
284        let mut got = Vec::new();
285        while let Ok(n) = consumer.read(&mut buf) {
286            if n == 0 {
287                break;
288            }
289            got.extend_from_slice(&buf[..n]);
290            if got.len() >= 27 * 4 + 19 + 27 + 4 + 259 {
291                break;
292            }
293        }
294        let mut pages = Vec::new();
295        let mut laces = Vec::new();
296        let mut at = 0;
297        while at + 27 <= got.len() {
298            assert_eq!(&got[at..at + 4], b"OggS");
299            let segments = got[at + 26] as usize;
300            let body: usize = got[at + 27..at + 27 + segments]
301                .iter()
302                .map(|&l| l as usize)
303                .sum();
304            let end = at + 27 + segments + body;
305            let mut copy = got[at..end].to_vec();
306            let stored = u32::from_le_bytes(copy[22..26].try_into().unwrap());
307            copy[22..26].copy_from_slice(&[0; 4]);
308            assert_eq!(crc32(&copy), stored, "page {} checksum", pages.len());
309            let granule = u64::from_le_bytes(copy[6..14].try_into().unwrap());
310            laces.push(copy[27..27 + segments].to_vec());
311            pages.push((granule, copy[27 + segments..].to_vec()));
312            at = end;
313        }
314        assert_eq!(pages.len(), 4);
315        assert!(pages[0].1.starts_with(b"OpusHead"));
316        assert_eq!(pages[0].1[9], 2);
317        assert!(pages[1].1.starts_with(b"OpusTags"));
318        assert_eq!(pages[2], (960, vec![0xfc, 1, 2, 3]));
319        assert_eq!(pages[3].0, 1920);
320        assert_eq!(pages[3].1.len(), 255);
321        assert_eq!(laces[3], vec![255, 0]);
322        drop(sink);
323        assert!(!std::path::Path::new(&path).exists());
324    }
325}