swarm-stt: only files that start like an audio container reach ffmpeg

docs/4-implementation/swarm/stt.md

@@ -31,7 +31,9 @@ into place with a rename, so a reader never sees half a transcript.
 The `Engine` trait takes bytes and returns text. `Whisper` is the real one:
 `ffmpeg` decodes any container to 16 kHz mono WAV, cutting it at `max_secs`,
 and `whisper-cli` from whisper.cpp reads the WAV with the configured model and
-language detection on. Decoding is pipe-only (`-protocol_whitelist pipe`): the audio comes from
+language detection on. Before anything decodes, `sniff` checks that the file starts like an audio
+container (Ogg, WAV, FLAC, MP3, AMR, MP4 or M4A, WebM); the daemon refuses a playlist, an image or a
+script without starting `ffmpeg`. Decoding is pipe-only (`-protocol_whitelist pipe`): the audio comes from
 rooms with outside people, so a file that names other files or URLs gets no
 way to open them. The service kills a run that outlasts its timeout (decode and transcribe
 together) and fails the request, so one hung child can't hold the single

swarm-stt/src/engine.rs

@@ -36,6 +36,12 @@ pub struct Whisper {
 #[async_trait]
 impl Engine for Whisper {
     async fn transcribe(&self, audio: &[u8]) -> Result<String, String> {
+        if !crate::sniff::is_audio(audio) {
+            return Err(
+                "the file does not look like audio (Ogg, WAV, FLAC, MP3, AMR, MP4/M4A, WebM)"
+                    .to_owned(),
+            );
+        }
         let dir = tempfile_in(&self.scratch)
             .await
             .map_err(|e| format!("making a scratch directory: {e}"))?;
@@ -157,8 +163,8 @@ mod tests {
     async fn audio_goes_through_ffmpeg_then_whisper_and_the_text_is_trimmed() {
         let dir = tempfile::tempdir().unwrap();
         let engine = whisper(dir.path(), FFMPEG_OK, WHISPER_ECHO);
-        let text = engine.transcribe(b"  hello there \n").await.unwrap();
-        assert_eq!(text, "hello there");
+        let text = engine.transcribe(b"OggS  hello there \n").await.unwrap();
+        assert_eq!(text, "OggS  hello there");
         assert_eq!(
             std::fs::read_dir(dir.path().join("scratch"))
                 .unwrap()
@@ -173,7 +179,7 @@ mod tests {
         let dir = tempfile::tempdir().unwrap();
         let mut engine = whisper(dir.path(), "sleep 60", WHISPER_ECHO);
         engine.timeout = std::time::Duration::from_millis(200);
-        let err = engine.transcribe(b"x").await.unwrap_err();
+        let err = engine.transcribe(b"OggSx").await.unwrap_err();
         assert!(err.contains("timed out"), "{err}");
         assert_eq!(
             std::fs::read_dir(dir.path().join("scratch"))
@@ -192,7 +198,7 @@ mod tests {
             record.display()
         );
         let engine = whisper(dir.path(), &ffmpeg, WHISPER_ECHO);
-        engine.transcribe(b"x").await.unwrap();
+        engine.transcribe(b"OggSx").await.unwrap();
         let args = std::fs::read_to_string(record).unwrap();
         assert!(
             args.contains("-protocol_whitelist pipe -i pipe:0"),
@@ -200,11 +206,25 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn a_file_that_is_not_audio_never_reaches_ffmpeg() {
+        let dir = tempfile::tempdir().unwrap();
+        let record = dir.path().join("ran");
+        let ffmpeg = format!("touch {}", record.display());
+        let engine = whisper(dir.path(), &ffmpeg, WHISPER_ECHO);
+        let err = engine
+            .transcribe(b"#EXTM3U\nfile:///etc/passwd")
+            .await
+            .unwrap_err();
+        assert!(err.contains("does not look like audio"), "{err}");
+        assert!(!record.exists(), "ffmpeg was started");
+    }
+
     #[tokio::test]
     async fn an_ffmpeg_failure_is_reported_with_its_stderr() {
         let dir = tempfile::tempdir().unwrap();
         let engine = whisper(dir.path(), "echo 'bad data' >&2; exit 1", WHISPER_ECHO);
-        let err = engine.transcribe(b"x").await.unwrap_err();
+        let err = engine.transcribe(b"OggSx").await.unwrap_err();
         assert!(
             err.contains("could not decode") && err.contains("bad data"),
             "{err}"
@@ -215,7 +235,7 @@ mod tests {
     async fn a_whisper_failure_is_reported_with_its_stderr() {
         let dir = tempfile::tempdir().unwrap();
         let engine = whisper(dir.path(), FFMPEG_OK, "echo 'no model' >&2; exit 2");
-        let err = engine.transcribe(b"x").await.unwrap_err();
+        let err = engine.transcribe(b"OggSx").await.unwrap_err();
         assert!(
             err.contains("whisper-cli failed") && err.contains("no model"),
             "{err}"

swarm-stt/src/lib.rs

@@ -9,12 +9,16 @@
 //! - an in-flight map, so two requests for one memo that arrive together share
 //!   one engine run instead of starting two.
 //!
+//! The audio comes from rooms with outside people, so [`sniff`] checks that a file
+//! starts like an audio container before any decoder sees it.
+//!
 //! The engine is behind the [`Engine`] trait. [`engine::Whisper`] decodes with
 //! `ffmpeg` and transcribes with `whisper-cli`; the tests use a fake.
 
 pub mod cache;
 pub mod engine;
 mod service;
+pub mod sniff;
 
 pub use engine::Engine;
 pub use service::{Config, Error, Reply, Service};

swarm-stt/src/sniff.rs

@@ -0,0 +1,64 @@
+//! A check of the first bytes of a file against the audio containers voice
+//! memos come in, before anything decodes it.
+//!
+//! The audio comes from matrix rooms with outside people, so the decoder only
+//! ever sees files that start like an audio container. A file that claims to
+//! be anything else (a playlist, an image, a script) stops here, which keeps
+//! most of `ffmpeg`'s demuxers out of reach.
+
+/// `true` when `bytes` start like Ogg (Opus, Vorbis), WAV, FLAC, MP3 (with or
+/// without an `ID3` tag), AMR, `MP4` or `M4A`, Matroska or `WebM`.
+#[must_use]
+pub fn is_audio(bytes: &[u8]) -> bool {
+    let starts = |magic: &[u8]| bytes.starts_with(magic);
+    starts(b"OggS")
+        || (starts(b"RIFF") && bytes.get(8..12) == Some(b"WAVE"))
+        || starts(b"fLaC")
+        || starts(b"ID3")
+        || starts(b"#!AMR")
+        // An MPEG audio frame: eleven set bits of sync.
+        || matches!(bytes, [0xFF, second, ..] if second & 0xE0 == 0xE0)
+        // The `ftyp` box of an ISO base media file (MP4, M4A, 3GP).
+        || bytes.get(4..8) == Some(b"ftyp")
+        // The EBML header of Matroska and WebM.
+        || starts(&[0x1A, 0x45, 0xDF, 0xA3])
+}
+
+#[cfg(test)]
+mod tests {
+    use super::is_audio;
+
+    #[test]
+    fn the_containers_voice_memos_come_in_pass() {
+        for (name, head) in [
+            ("ogg", &b"OggS\0\x02"[..]),
+            ("wav", b"RIFF\x24\0\0\0WAVEfmt "),
+            ("flac", b"fLaC\0"),
+            ("mp3 with a tag", b"ID3\x04\0"),
+            ("mp3 frame", &[0xFF, 0xFB, 0x90, 0x00]),
+            ("amr", b"#!AMR\n"),
+            ("m4a", b"\0\0\0\x20ftypM4A "),
+            ("webm", &[0x1A, 0x45, 0xDF, 0xA3, 0x9F]),
+        ] {
+            assert!(is_audio(head), "{name} was refused");
+        }
+    }
+
+    #[test]
+    fn files_that_are_not_audio_containers_are_refused() {
+        for (name, head) in [
+            ("an HLS playlist", &b"#EXTM3U\n#EXT-X-VERSION:3\n"[..]),
+            ("a png", b"\x89PNG\r\n\x1a\n"),
+            ("a script", b"#!/bin/sh\nid\n"),
+            (
+                "a concat list",
+                b"ffconcat version 1.0\nfile '/etc/passwd'\n",
+            ),
+            ("a wave-less riff (avi)", b"RIFF\x24\0\0\0AVI LIST"),
+            ("empty", b""),
+            ("one byte", &[0xFF]),
+        ] {
+            assert!(!is_audio(head), "{name} was accepted");
+        }
+    }
+}