tts: bound the audio cache and request length

1 files changed, 86 insertions, 5 deletions

diff --git a/src/tts.rs b/src/tts.rs
index 3ae5140..dbe0d88 100644
--- a/src/tts.rs
+++ b/src/tts.rs
@@ -1,6 +1,7 @@
 use std::io::Cursor;
 use std::path::PathBuf;
 use std::sync::OnceLock;
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::time::{Duration, Instant};
 
 use dashmap::DashMap;
@@ -18,6 +19,14 @@ const TARGET_SAMPLE_RATE: u32 = 48_000;
 /// Cache entries older than this are evicted on access.
 const CACHE_TTL: Duration = Duration::from_secs(3600);
 
+/// Hard cap on cached audio so arbitrary `/tts` text cannot exhaust memory.
+const CACHE_MAX_ENTRIES: usize = 256;
+const CACHE_MAX_BYTES: usize = 64 * 1024 * 1024;
+
+/// Longest text accepted for synthesis. Keeps a single request (and its
+/// cached PCM) bounded.
+const MAX_TTS_CHARS: usize = 2_000;
+
 /// All voice-state TTS event types.
 const EVENT_TYPES: &[&str] = &[
     "joined",
@@ -37,6 +46,7 @@ const EVENT_TYPES: &[&str] = &[
 pub struct TtsState {
     piper: OnceLock<Mutex<Piper>>,
     cache: DashMap<String, (Vec<u8>, Instant)>,
+    cache_bytes: AtomicUsize,
     model_path: Option<PathBuf>,
 }
 
@@ -45,6 +55,7 @@ impl TtsState {
         Self {
             piper: OnceLock::new(),
             cache: DashMap::new(),
+            cache_bytes: AtomicUsize::new(0),
             model_path,
         }
     }
@@ -93,22 +104,65 @@ impl TtsState {
         let entry = self.cache.get(text)?;
         if entry.1.elapsed() > CACHE_TTL {
             drop(entry);
-            self.cache.remove(text);
+            if let Some((old, _)) = self.cache.remove(text) {
+                self.cache_bytes.fetch_sub(old.len(), Ordering::Relaxed);
+            }
             None
         } else {
             Some(entry.0.clone())
         }
     }
 
-    /// Store audio bytes in the cache.
+    /// Store audio bytes in the cache, evicting expired/old entries first.
     fn cache_store(&self, text: String, audio: Vec<u8>) {
-        self.cache.insert(text, (audio, Instant::now()));
+        let audio_len = audio.len();
+        if audio_len > CACHE_MAX_BYTES {
+            return;
+        }
+        self.cache_evict();
+        while self.cache.len() >= CACHE_MAX_ENTRIES {
+            if !self.cache_evict_oldest() {
+                break;
+            }
+        }
+        while self.cache_bytes.load(Ordering::Relaxed) + audio_len > CACHE_MAX_BYTES {
+            if !self.cache_evict_oldest() {
+                break;
+            }
+        }
+
+        if let Some((old, _)) = self.cache.insert(text, (audio, Instant::now())) {
+            self.cache_bytes.fetch_sub(old.len(), Ordering::Relaxed);
+        }
+        self.cache_bytes.fetch_add(audio_len, Ordering::Relaxed);
     }
 
     /// Evict all expired cache entries.
     fn cache_evict(&self) {
-        self.cache
-            .retain(|_, (_, inserted_at)| inserted_at.elapsed() <= CACHE_TTL);
+        let bytes = &self.cache_bytes;
+        self.cache.retain(|_, (audio, inserted_at)| {
+            let keep = inserted_at.elapsed() <= CACHE_TTL;
+            if !keep {
+                bytes.fetch_sub(audio.len(), Ordering::Relaxed);
+            }
+            keep
+        });
+    }
+
+    /// Evict the least-recently-inserted entry. Returns false when empty.
+    fn cache_evict_oldest(&self) -> bool {
+        let oldest = self
+            .cache
+            .iter()
+            .min_by_key(|entry| entry.value().1)
+            .map(|entry| entry.key().clone());
+        let Some(key) = oldest else {
+            return false;
+        };
+        if let Some((old, _)) = self.cache.remove(&key) {
+            self.cache_bytes.fetch_sub(old.len(), Ordering::Relaxed);
+        }
+        true
     }
 
     /// Synthesize text → resampled 48 kHz stereo PCM bytes.
@@ -116,6 +170,11 @@ impl TtsState {
     /// Checks cache first. On miss, runs piper inference, resamples, caches,
     /// and returns.
     pub async fn synthesize(&self, text: &str) -> Result<Vec<u8>> {
+        if text.chars().count() > MAX_TTS_CHARS {
+            return Err(Error::Custom(format!(
+                "TTS text is too long (max {MAX_TTS_CHARS} characters)"
+            )));
+        }
         // Cache hit
         if let Some(cached) = self.cache_lookup(text) {
             debug!(text_len = text.len(), "tts cache hit");
@@ -289,3 +348,25 @@ pub fn pcm_bytes_to_input(pcm_f32_bytes: Vec<u8>) -> songbird::input::Input {
     // 2 channels (stereo), 48 kHz — matches TARGET_SAMPLE_RATE
     RawAdapter::new(cursor, TARGET_SAMPLE_RATE, 2).into()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn cache_is_bounded_by_entries_and_bytes() {
+        let state = TtsState::new(None);
+        for i in 0..(CACHE_MAX_ENTRIES + 32) {
+            state.cache_store(format!("phrase-{i}"), vec![0u8; 1024]);
+        }
+        assert!(state.cache.len() <= CACHE_MAX_ENTRIES);
+        assert!(state.cache_bytes.load(Ordering::Relaxed) <= CACHE_MAX_BYTES);
+    }
+
+    #[tokio::test]
+    async fn overlong_text_is_rejected() {
+        let state = TtsState::new(None);
+        let text = "x".repeat(MAX_TTS_CHARS + 1);
+        assert!(state.synthesize(&text).await.is_err());
+    }
+}