fix(companion): preserve uncached caption responses

spec/invidious/subtitle_cache_spec.cr

@@ -11,6 +11,13 @@ require "../../src/invidious/config"
 
 CONFIG.invidious_companion_key = "1234567890123456"
 
+private def subtitle_response(status_code : Int32, content_type : String, body : String)
+  headers = HTTP::Headers.new
+  headers["Content-Type"] = content_type
+  headers["Access-Control-Allow-Origin"] = "*"
+  Invidious::SubtitleCache::Response.new(status_code, headers, body)
+end
+
 Spectator.describe Invidious::SubtitleCache do
   describe "check verification" do
     it "verifies valid tokens" do
@@ -77,22 +84,22 @@ Spectator.describe Invidious::SubtitleCache do
       fetch_count = 0
       vtt = "WEBVTT\n\n00:00:00.000 --> 00:00:01.000\nHello"
 
-      entry1, status1 = cache.get_or_fetch("video:abc") do
+      result1 = cache.get_or_fetch("video:abc") do
         fetch_count += 1
-        {200, "text/vtt", vtt}
+        subtitle_response(200, "text/vtt", vtt)
       end
 
-      expect(status1).to eq("miss")
+      expect(result1.cache_status).to eq("miss")
-      expect(entry1.try &.body).to eq(vtt)
+      expect(result1.entry.try &.body).to eq(vtt)
       expect(fetch_count).to eq(1)
 
-      entry2, status2 = cache.get_or_fetch("video:abc") do
+      result2 = cache.get_or_fetch("video:abc") do
         fetch_count += 1
-        {200, "text/vtt", vtt}
+        subtitle_response(200, "text/vtt", vtt)
       end
 
-      expect(status2).to eq("hit")
+      expect(result2.cache_status).to eq("hit")
-      expect(entry2.try &.body).to eq(vtt)
+      expect(result2.entry.try &.body).to eq(vtt)
       expect(fetch_count).to eq(1)
     end
 
@@ -104,12 +111,12 @@ Spectator.describe Invidious::SubtitleCache do
 
       5.times do
         spawn do
-          entry, status = cache.get_or_fetch("video:concurrent") do
+          result = cache.get_or_fetch("video:concurrent") do
             fetch_count += 1
             sleep 0.05.seconds
-            {200, "text/vtt", vtt}
+            subtitle_response(200, "text/vtt", vtt)
           end
-          done_ch.send(status)
+          done_ch.send(result.cache_status)
         end
       end
 
@@ -159,12 +166,28 @@ Spectator.describe Invidious::SubtitleCache do
 
     it "does not cache failed or non-vtt responses" do
       cache = Invidious::SubtitleCache.new
-      entry, status = cache.get_or_fetch("video:bad") do
+      result = cache.get_or_fetch("video:bad") do
-        {404, "application/json", "{\"error\":\"not found\"}"}
+        subtitle_response(404, "application/json", "{\"error\":\"not found\"}")
       end
 
-      expect(status).to eq("bypass")
+      expect(result.cache_status).to eq("bypass")
-      expect(entry).to be_nil
+      expect(result.entry).to be_nil
+      expect(result.response.try &.body).to eq("{\"error\":\"not found\"}")
+      expect(result.response.try &.headers["Access-Control-Allow-Origin"]).to eq("*")
+      expect(cache.size).to eq(0)
+    end
+
+    it "does not cache oversized fetch responses" do
+      cache = Invidious::SubtitleCache.new
+      oversized = "WEBVTT\n" + ("x" * (Invidious::SubtitleCache::MAX_ENTRY_BYTES + 1))
+
+      result = cache.get_or_fetch("video:oversized") do
+        subtitle_response(200, "text/vtt", oversized)
+      end
+
+      expect(result.cache_status).to eq("bypass")
+      expect(result.entry).to be_nil
+      expect(result.response.try &.body).to eq(oversized)
       expect(cache.size).to eq(0)
     end
   end

src/invidious/helpers/subtitle_cache.cr

@@ -15,6 +15,24 @@ module Invidious
       end
     end
 
+    struct Response
+      getter status_code : Int32
+      getter headers : HTTP::Headers
+      getter body : String
+
+      def initialize(@status_code : Int32, @headers : HTTP::Headers, @body : String)
+      end
+    end
+
+    struct FetchResult
+      getter entry : Entry?
+      getter response : Response?
+      getter cache_status : String
+
+      def initialize(@entry : Entry?, @response : Response?, @cache_status : String)
+      end
+    end
+
     DEFAULT_MAX_ENTRIES = 256
     DEFAULT_MAX_BYTES   = 64 * 1024 * 1024 # 64 MiB
     DEFAULT_TTL         = 6.hours          # 21600 seconds
@@ -31,7 +49,7 @@ module Invidious
       @entries = Hash(String, Entry).new
       @total_bytes = 0
       @mutex = Mutex.new
-      @in_flight = Hash(String, Array(::Channel(Entry?))).new
+      @in_flight = Hash(String, Array(::Channel(FetchResult))).new
     end
 
     def self.valid_vtt?(body : String) : Bool
@@ -75,15 +93,15 @@ module Invidious
     end
 
     def put(key : String, body : String, content_type : String) : Entry?
-      return nil unless SubtitleCache.valid_vtt?(body)
-      return nil if body.bytesize > MAX_ENTRY_BYTES
-
       @mutex.synchronize do
         put_internal(key, body, content_type)
       end
     end
 
     private def put_internal(key : String, body : String, content_type : String) : Entry?
+      return nil unless SubtitleCache.valid_vtt?(body)
+      return nil if body.bytesize > MAX_ENTRY_BYTES || body.bytesize > @max_bytes
+
       if old = @entries.delete(key)
         @total_bytes -= old.size
       end
@@ -104,48 +122,40 @@ module Invidious
       entry
     end
 
-    def get_or_fetch(key : String, &fetch_block : -> Tuple(Int32, String, String)?) : Tuple(Entry?, String)
+    def get_or_fetch(key : String, &fetch_block : -> Response?) : FetchResult
-      wait_ch : ::Channel(Entry?)? = nil
+      wait_ch : ::Channel(FetchResult)? = nil
 
       @mutex.synchronize do
         if entry = get_internal(key)
-          return {entry, "hit"}
+          return FetchResult.new(entry, nil, "hit")
         end
 
         if waiting_list = @in_flight[key]?
-          ch = ::Channel(Entry?).new(1)
+          ch = ::Channel(FetchResult).new(1)
           waiting_list << ch
           wait_ch = ch
         else
-          @in_flight[key] = Array(::Channel(Entry?)).new
+          @in_flight[key] = Array(::Channel(FetchResult)).new
         end
       end
 
       if ch = wait_ch
-        entry = ch.receive
+        return ch.receive
-        if entry
-          return {entry, "hit"}
-        else
-          return {nil, "bypass"}
-        end
       end
 
       # Primary fetcher for this key
-      fetch_result = begin
+      response = begin
         fetch_block.call
       rescue
         nil
       end
 
       cached_entry : Entry? = nil
-      waiting_channels = Array(::Channel(Entry?)).new
+      waiting_channels = Array(::Channel(FetchResult)).new
 
       @mutex.synchronize do
-        if fetch_result
+        if response && response.status_code == 200
-          status, content_type, body = fetch_result
+          cached_entry = put_internal(key, response.body, response.headers["Content-Type"]? || "text/vtt; charset=utf-8")
-          if status == 200 && SubtitleCache.valid_vtt?(body)
-            cached_entry = put_internal(key, body, content_type)
-          end
         end
 
         if channels = @in_flight.delete(key)
@@ -153,15 +163,14 @@ module Invidious
         end
       end
 
+      primary_result = FetchResult.new(cached_entry, response, cached_entry ? "miss" : "bypass")
+      waiter_result = FetchResult.new(cached_entry, response, cached_entry ? "hit" : "bypass")
+
       waiting_channels.each do |w_ch|
-        w_ch.send(cached_entry)
+        w_ch.send(waiter_result)
       end
 
-      if cached_entry
+      primary_result
-        {cached_entry, "miss"}
-      else
-        {nil, "bypass"}
-      end
     end
   end
 end

src/invidious/routes/companion.cr

@@ -82,36 +82,37 @@ module Invidious::Routes::Companion
 
     cache_key = "video:#{video_id}|label:#{label}|lang:#{lang}|tlang:#{tlang}"
 
-    entry, cache_status = SUBTITLE_CACHE.get_or_fetch(cache_key) do
+    result = SUBTITLE_CACHE.get_or_fetch(cache_key) do
       fetch_from_companion(url, env.request.headers)
     end
 
-    if entry
+    if entry = result.entry
       env.response.status_code = 200
+      env.response.headers["Access-Control-Allow-Origin"] = "*"
       env.response.headers["Content-Type"] = entry.content_type
       env.response.headers["Cache-Control"] = "private, max-age=21600"
-      env.response.headers["X-Invidious-Subtitle-Cache"] = cache_status
+      env.response.headers["X-Invidious-Subtitle-Cache"] = result.cache_status
       env.response.print entry.body
       return
-    else
+    elsif response = result.response
-      begin
+      env.response.status_code = response.status_code
-        COMPANION_POOL.client do |wrapper|
+      response.headers.each do |key, value|
-          wrapper.client.get(url, env.request.headers) do |resp|
+        env.response.headers[key] = value
-            env.response.headers["X-Invidious-Subtitle-Cache"] = cache_status
-            return self.proxy_companion(env, resp)
-          end
-        end
-      rescue ex
       end
+      env.response.headers["X-Invidious-Subtitle-Cache"] = result.cache_status
+      env.response.print response.body
     end
   end
 
-  private def self.fetch_from_companion(url : String, headers : HTTP::Headers) : Tuple(Int32, String, String)?
+  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 = resp.body_io.gets_to_end
-        content_type = resp.headers["Content-Type"]? || "text/vtt; charset=utf-8"
+        response_headers = HTTP::Headers.new
-        return {resp.status_code, content_type, body}
+        resp.headers.each do |key, value|
+          response_headers[key] = value
+        end
+        return Invidious::SubtitleCache::Response.new(resp.status_code, response_headers, body)
       end
     end
   rescue ex