fix(companion): handle caption cache edge cases

spec/invidious/subtitle_cache_spec.cr

@@ -149,6 +149,26 @@ Spectator.describe Invidious::SubtitleCache do
       expect(cache.get("k2")).to be_nil
     end
 
+    it "rejects entries when entry capacity is zero" do
+      cache = Invidious::SubtitleCache.new(max_entries: 0)
+      vtt = "WEBVTT\n\n00:00:00.000 --> 00:00:01.000\nZero capacity"
+
+      expect(cache.put("k", vtt, "text/vtt")).to be_nil
+      expect(cache.size).to eq(0)
+
+      byte_limited = Invidious::SubtitleCache.new(max_bytes: 0)
+      expect(byte_limited.put("k", vtt, "text/vtt")).to be_nil
+      expect(byte_limited.size).to eq(0)
+    end
+
+    it "matches exact caption proxy paths only" do
+      regex = Invidious::SubtitleCache::CAPTION_PATH_REGEX
+
+      expect(regex.match("/api/v1/captions/video_id").try &.[1]).to eq("video_id")
+      expect(regex.match("/companion/api/v1/captions/video_id").try &.[1]).to eq("video_id")
+      expect(regex.match("/api/v1/captions/video_id/extra")).to be_nil
+    end
+
     it "evicts oldest entries when max_bytes is exceeded" do
       cache = Invidious::SubtitleCache.new(max_entries: 10, max_bytes: 60)
       vtt1 = "WEBVTT\n\n12345678901234567890" # ~30 bytes
@@ -195,10 +215,14 @@ Spectator.describe Invidious::SubtitleCache do
 
     it "stops reading bodies beyond the entry size limit" do
       valid_body = "WEBVTT\n"
-      expect(Invidious::SubtitleCache.read_limited_body(IO::Memory.new(valid_body))).to eq(valid_body)
+      valid_result = Invidious::SubtitleCache.read_limited_body(IO::Memory.new(valid_body))
+      expect(valid_result.oversized).to be_false
+      expect(valid_result.body).to eq(valid_body)
 
       oversized_body = "x" * (Invidious::SubtitleCache::MAX_ENTRY_BYTES + 1)
-      expect(Invidious::SubtitleCache.read_limited_body(IO::Memory.new(oversized_body))).to be_nil
+      oversized_result = Invidious::SubtitleCache.read_limited_body(IO::Memory.new(oversized_body))
+      expect(oversized_result.oversized).to be_true
+      expect(oversized_result.body).to be_nil
     end
   end
 end

src/invidious/helpers/subtitle_cache.cr

@@ -1,5 +1,7 @@
 module Invidious
   class SubtitleCache
+    CAPTION_PATH_REGEX = %r{^/(?:companion/)?api/v1/captions/([^/?#]+)$}
+
     struct Entry
       getter body : String
       getter content_type : String
@@ -33,6 +35,14 @@ module Invidious
       end
     end
 
+    struct LimitedBody
+      getter body : String?
+      getter oversized : Bool
+
+      def initialize(@body : String?, @oversized : Bool)
+      end
+    end
+
     DEFAULT_MAX_ENTRIES = 256
     DEFAULT_MAX_BYTES   = 64 * 1024 * 1024 # 64 MiB
     DEFAULT_TTL         = 6.hours          # 21600 seconds
@@ -62,7 +72,7 @@ module Invidious
       header_end == trimmed.size || " \t\r\n".includes?(trimmed[header_end])
     end
 
-    def self.read_limited_body(input : IO, limit : Int32 = MAX_ENTRY_BYTES) : String?
+    def self.read_limited_body(input : IO, limit : Int32 = MAX_ENTRY_BYTES) : LimitedBody
       output = IO::Memory.new
       buffer = Bytes.new(READ_CHUNK_BYTES)
 
@@ -74,8 +84,8 @@ module Invidious
         output.write(buffer[0, bytes_read])
       end
 
-      return nil if input.read_byte
-      output.to_s
+      oversized = !input.read_byte.nil?
+      LimitedBody.new(oversized ? nil : output.to_s, oversized)
     end
 
     def size : Int32
@@ -120,6 +130,7 @@ module Invidious
 
     private def put_internal(key : String, body : String, content_type : String) : Entry?
       return nil unless SubtitleCache.valid_vtt?(body)
+      return nil if @max_entries <= 0 || @max_bytes <= 0
       return nil if body.bytesize > MAX_ENTRY_BYTES || body.bytesize > @max_bytes
 
       if old = @entries.delete(key)

src/invidious/routes/companion.cr

@@ -9,7 +9,7 @@ module Invidious::Routes::Companion
     end
 
     path = env.request.path.rstrip('/')
-    if match = path.match(%r{^/(?:companion/)?api/v1/captions/([^/?#]+)})
+    if match = path.match(Invidious::SubtitleCache::CAPTION_PATH_REGEX)
       video_id = match[1]
       return self.handle_caption_request(env, url, video_id)
     end
@@ -107,7 +107,13 @@ module Invidious::Routes::Companion
   private def self.fetch_from_companion(url : String, headers : HTTP::Headers) : Invidious::SubtitleCache::Response?
     COMPANION_POOL.client do |wrapper|
       wrapper.client.get(url, headers) do |resp|
-        body = Invidious::SubtitleCache.read_limited_body(resp.body_io)
+        limited_body = Invidious::SubtitleCache.read_limited_body(resp.body_io)
+        if limited_body.oversized
+          resp.body_io.close
+          return oversized_caption_response
+        end
+
+        body = limited_body.body
         return nil unless body
         response_headers = HTTP::Headers.new
         resp.headers.each do |key, value|
@@ -120,6 +126,14 @@ module Invidious::Routes::Companion
     nil
   end
 
+  private def self.oversized_caption_response
+    headers = HTTP::Headers.new
+    headers["Access-Control-Allow-Origin"] = "*"
+    headers["Cache-Control"] = "no-store"
+    headers["Content-Type"] = "text/plain; charset=utf-8"
+    Invidious::SubtitleCache::Response.new(413, headers, "Caption response exceeds size limit\n")
+  end
+
   private def self.proxy_companion(env, response)
     env.response.status_code = response.status_code
     response.headers.each do |key, value|