swarm-stt: time-box a run and decode pipe-only

swarm-stt/src/engine.rs

@@ -28,6 +28,9 @@ pub struct Whisper {
     pub scratch: PathBuf,
     /// Seconds of audio to decode; the rest is cut.
     pub max_secs: u64,
+    /// How long one run (decode and transcribe together) may take. The service
+    /// runs one job at a time, so a hung child would block every request.
+    pub timeout: std::time::Duration,
 }
 
 #[async_trait]
@@ -36,7 +39,11 @@ impl Engine for Whisper {
         let dir = tempfile_in(&self.scratch)
             .await
             .map_err(|e| format!("making a scratch directory: {e}"))?;
-        let result = self.run(&dir, audio).await;
+        let result = match tokio::time::timeout(self.timeout, self.run(&dir, audio)).await {
+            Ok(result) => result,
+            // Dropping `run` kills the child that was running.
+            Err(_) => Err(format!("timed out after {}s", self.timeout.as_secs())),
+        };
         let _ = tokio::fs::remove_dir_all(&dir).await;
         result
     }
@@ -46,7 +53,19 @@ impl Whisper {
     async fn run(&self, dir: &std::path::Path, audio: &[u8]) -> Result<String, String> {
         let wav = dir.join("audio.wav");
         let mut decode = Command::new(&self.ffmpeg)
-            .args(["-nostdin", "-v", "error", "-i", "pipe:0", "-t"])
+            // Pipe-only: the audio is untrusted, and a demuxer that follows
+            // references inside it (an HLS playlist, say) must not open files
+            // or sockets.
+            .args([
+                "-nostdin",
+                "-v",
+                "error",
+                "-protocol_whitelist",
+                "pipe",
+                "-i",
+                "pipe:0",
+                "-t",
+            ])
             .arg(self.max_secs.to_string())
             .args(["-ar", "16000", "-ac", "1", "-y"])
             .arg(&wav)
@@ -124,6 +143,7 @@ mod tests {
             model: dir.join("model.bin"),
             scratch: dir.join("scratch"),
             max_secs: 300,
+            timeout: std::time::Duration::from_secs(5),
         }
     }
 
@@ -148,6 +168,38 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn a_decode_that_never_ends_times_out_and_leaves_no_scratch_directory() {
+        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();
+        assert!(err.contains("timed out"), "{err}");
+        assert_eq!(
+            std::fs::read_dir(dir.path().join("scratch"))
+                .unwrap()
+                .count(),
+            0
+        );
+    }
+
+    #[tokio::test]
+    async fn ffmpeg_is_started_pipe_only() {
+        let dir = tempfile::tempdir().unwrap();
+        let record = dir.path().join("args");
+        let ffmpeg = format!(
+            r#"echo "$@" > {}; for last; do :; done; cat > "$last""#,
+            record.display()
+        );
+        let engine = whisper(dir.path(), &ffmpeg, WHISPER_ECHO);
+        engine.transcribe(b"x").await.unwrap();
+        let args = std::fs::read_to_string(record).unwrap();
+        assert!(
+            args.contains("-protocol_whitelist pipe -i pipe:0"),
+            "{args}"
+        );
+    }
+
     #[tokio::test]
     async fn an_ffmpeg_failure_is_reported_with_its_stderr() {
         let dir = tempfile::tempdir().unwrap();